package cn.jbolt.ai.core.chain.node.provider;

import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.app.config.AppConfigConstant;
import cn.jbolt.ai.app.config.ChatNodeConfig;
import cn.jbolt.ai.app.config.ChatWithKlbNodeConfig;
import cn.jbolt.ai.app.entity.AiAppMessage;
import cn.jbolt.ai.app.service.AiAppMessageService;
import cn.jbolt.ai.app.service.AiAppService;
import cn.jbolt.ai.core.chain.ChainContext;
import cn.jbolt.ai.klb.FileType;
import cn.jbolt.ai.klb.entity.AiDoc;
import cn.jbolt.util.JSONUtil;
import com.fasterxml.jackson.databind.JsonNode;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.capability.model.JBoltAI;
import com.jboltai.event.Event;
import com.jboltai.event.ai.aiknowledge.AIKnowledgeEvent;
import com.jboltai.event.signal.SuccessEvent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;


/**
 * 知识库回答的节点
 */
@Component
public class ChatWithKlbNodeProvider implements EventNodeProvider<Event> {

    @Autowired
    private AiAppService aiAppService;
    @Autowired
    private AiAppMessageService aiAppMessageService;
    private static final Logger logger = LoggerFactory.getLogger(ChatWithKlbNodeProvider.class);

    public Event provide(JsonNode nodeConfig, ChainContext context) {
        try {
            ChatNodeConfig llmConfig = ChatNodeConfig.from(nodeConfig.get(AppConfigConstant.CONFIG_TYPE_LLM));
            ChatWithKlbNodeConfig chatWithKlbConfig = ChatWithKlbNodeConfig.from(nodeConfig.get(AppConfigConstant.CONFIG_TYPE_CHAT_WITH_KLB));
            if (chatWithKlbConfig == null) {
                logger.error("APP [{}:{}] 节点未配置", context.getAiApp().getSn() ,context.getAiApp().getName());
                throw new RuntimeException("节点未配置" );
            }
            List<AiDoc> docs = context.getAttr(chatWithKlbConfig.getKlbResultField());

            if ((docs == null || docs.size() == 0) && chatWithKlbConfig.isOnlyKlb() ) {
                //仅知识库模式，还没查到相关资料
                if (llmConfig.isOutputToResponse()) {
                    context.sendResponseMsg(chatWithKlbConfig.getNoMatchReply(), "markdown");
                }
                if (llmConfig.isPersistDb()) {
                    //把答案给存储
                    context.getResponseMsg().setContent(chatWithKlbConfig.getNoMatchReply());
                    context.getResponseMsg().setStatus(AiAppMessage.STATUS_SUCCESS);
                    aiAppMessageService.update(context.getResponseMsg());
                    return new SuccessEvent(null);
                }
            }

            //向前端推送参考
            if (chatWithKlbConfig.isShowDocRefers()) {
                context.sendReference(docs);
            }

            String question = context.getAttr(chatWithKlbConfig.getPromptField());
            List<String> background = generateBackgroundText(docs);
            AIKnowledgeEvent event = JBoltAI.chatWithKnowledge(question, background, chatWithKlbConfig.getNoMatchReply());
            event.setKnowledgeOnly(chatWithKlbConfig.isOnlyKlb());
            //根据配置初始化event
            configEvent(llmConfig, context, event);
            event.onSuccess((e, res) -> {
                if (llmConfig.isOutputToResponse()) {
                    //如果配置项要求 内容返回到响应里，那就输出，默认是true的
                    //构建相应信息
                    context.sendResponseMsg(res.getContent(), llmConfig.getFormat());
                }
            });
            event.onFail((e, error) -> {
                if (llmConfig.isOutputToResponse()) {
                    //如果配置项要求 内容返回到响应里，那就输出
                    //构建相应信息
                    context.sendErrorMsg(error.getMsg());

                }
            });

            event.onComplete((e) -> {
                if (e.isSuccess()) { //失败
                    //事件执行成功
                    if (StrUtil.isNotBlank(chatWithKlbConfig.getOutputField())) {
                        //配置项要求，把结果输出到上下文
                        context.setAttr(chatWithKlbConfig.getOutputField(), e.getResult());
                    }
                    if (llmConfig.isPersistDb()) {//需要持久化数据
                        //把答案给存储
                        context.getResponseMsg().setContent(e.getResult().getContent());
                        context.getResponseMsg().setStatus(AiAppMessage.STATUS_SUCCESS);
                        context.getResponseMsg().setDocs(JSONUtil.toJson(docs.stream().map(AiDoc::getId).collect(Collectors.toList())));
                        aiAppMessageService.update(context.getResponseMsg());
                    }
                } else { //失败
                    if (llmConfig.isPersistDb()) {//需要持久化数据
                        context.getResponseMsg().setStatus(e.isFail() ? AiAppMessage.STATUS_ERROR : AiAppMessage.STATUS_CANCEL);
                        context.getResponseMsg().setErrorMsg(e.getError() != null ? e.getError().getMsg(): "");
                        aiAppMessageService.update(context.getResponseMsg());
                    }

                }
            });
            return event;
        } catch (Exception e) {
            e.printStackTrace();
            logger.error("ChatNodeHandler process 异常,异常原因", e);
            context.sendErrorMsg( "网络异常,请稍后重试....");
        }
        return null;
    }

    private List<String> generateBackgroundText(List<AiDoc> docs) {
        if (docs == null || docs.size() == 0) {
            return Collections.emptyList();
        }
        //将结果保存到 输出字段中
        List<String> list = new ArrayList<>();
        for (AiDoc aiDoc : docs) {
            if (FileType.QA.getCode().equals(aiDoc.getDocType())) {
                list.add(aiDoc.getTitle() + "\n" + aiDoc.getAnswer());
            } else {
                list.add(aiDoc.getTitle());
            }
        }
        return list;
    }

    private void configEvent(ChatNodeConfig config, ChainContext context, AIKnowledgeEvent event) {

        event.setStream(config.isStream());
        if (config.getModelType() != null) {
            event.setModel(config.getModelType());
        }
        if (StrUtil.isNotBlank(config.getDefaultModel())) {
            event.setModelName(config.getDefaultModel());
        }
        if (config.getMaxTokens() > 0) {
            event.setMaxTokens(config.getMaxTokens());
        }
        if (StrUtil.isNotBlank(config.getSystemRole())) {
            event.systemMessage(config.getSystemRole());
        }
        if (config.isUseContext()) {
            Long appId = context.getAiApp().getId();
            Long sessionId = context.getAiSession().getId();
            List<AIMessage> messageList = aiAppService.getAppSessionContextMessages(sessionId, appId);
            if (messageList != null && messageList.size() > 0) {
                event.contextMessages(messageList);
            }
        }
        if (config.getTemperature() != null&& config.getTemperature()>0) {
            event.setTemperature(config.getTemperature());
        }
        if (config.getTopP() != null&& config.getTopP()>0) {
            event.setTopP(config.getTopP());
        }
        if (config.getTopK() != null&& config.getTopK()>0) {
            event.setTopK(config.getTopK());
        }

    }



}
